{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "75b58048-7d14-4fc6-8085-1fc08c81b4a6",
      "metadata": {
        "id": "75b58048-7d14-4fc6-8085-1fc08c81b4a6"
      },
      "source": [
        "# Learn OpenAI Whisper - Chapter 4\n",
        "## Fine-tune Whisper with Hugging Face transformers\n",
        "\n",
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1LADNomT0JUBCsopU6r_NsZfOaNiz2N3h)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "fbfa8ad5-4cdc-4512-9058-836cbbf65e1a",
      "metadata": {
        "id": "fbfa8ad5-4cdc-4512-9058-836cbbf65e1a"
      },
      "source": [
        "In this Colab, we present a step-by-step guide on how to fine-tune Whisper\n",
        "for any multilingual ASR dataset using Hugging Face 🤗 Transformers."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "55fb8d21-df06-472a-99dd-b59567be6dad",
      "metadata": {
        "id": "55fb8d21-df06-472a-99dd-b59567be6dad"
      },
      "source": [
        "## Prepare environment"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "844a4861-929c-4762-b29b-80b1e95aba4b",
      "metadata": {
        "id": "844a4861-929c-4762-b29b-80b1e95aba4b"
      },
      "source": [
        "First of all, let's try to secure a decent GPU for our Colab! Unfortunately, it's becoming much harder to get access to a good GPU with the free version of Google Colab. However, with Google Colab Pro one should have no issues in being allocated a V100 or P100 GPU.\n",
        "\n",
        "To get a GPU, click _Runtime_ -> _Change runtime type_, then change _Hardware accelerator_ from _None_ to _GPU_."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9abea5d7-9d54-434b-a6bd-399d1b3c6c1a",
      "metadata": {
        "id": "9abea5d7-9d54-434b-a6bd-399d1b3c6c1a"
      },
      "source": [
        "We can verify that we've been assigned a GPU and view its specifications:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "95048026-a3b7-43f0-a274-1bad65e407b4",
      "metadata": {
        "id": "95048026-a3b7-43f0-a274-1bad65e407b4"
      },
      "outputs": [],
      "source": [
        "gpu_info = !nvidia-smi\n",
        "gpu_info = '\\n'.join(gpu_info)\n",
        "if gpu_info.find('failed') >= 0:\n",
        "  print('Not connected to a GPU')\n",
        "else:\n",
        "  print(gpu_info)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "1d85d613-1c7e-46ac-9134-660bbe7ebc9d",
      "metadata": {
        "id": "1d85d613-1c7e-46ac-9134-660bbe7ebc9d"
      },
      "source": [
        "We'll employ several popular Python packages to fine-tune the Whisper model.\n",
        "We'll use `datasets` to download and prepare our training data and\n",
        "`transformers` to load and train our Whisper model. We'll also require\n",
        "the `soundfile` package to pre-process audio files, `evaluate` and `jiwer` to\n",
        "assess the performance of our model. Finally, we'll\n",
        "use `gradio` to build a flashy demo of our fine-tuned model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e68ea9f8-9b61-414e-8885-3033b67c2850",
      "metadata": {
        "id": "e68ea9f8-9b61-414e-8885-3033b67c2850"
      },
      "outputs": [],
      "source": [
        "!pip install -q datasets>=2.6.1\n",
        "!pip install -q git+https://github.com/huggingface/transformers\n",
        "!pip install -q librosa\n",
        "!pip install -q evaluate>=0.30\n",
        "!pip install -q jiwer\n",
        "!pip install -q gradio"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -q accelerate -U"
      ],
      "metadata": {
        "id": "J9wGvNnjmg4I"
      },
      "id": "J9wGvNnjmg4I",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "id": "1f60d173-8de1-4ed7-bc9a-d281cf237203",
      "metadata": {
        "id": "1f60d173-8de1-4ed7-bc9a-d281cf237203"
      },
      "source": [
        "We highly recommend directly uploading model checkpoints to the [Hugging Face Hub](https://huggingface.co/) while you train. The Hub offers several benefits:\n",
        "- **Version Control**: Ensures every model checkpoint is preserved throughout the training process.\n",
        "- **Tensorboard Logs**: Allows for the monitoring of key metrics as training progresses.\n",
        "- **Model Cards**: Provides a space to explain the model's functionality and its appropriate applications.\n",
        "- **Community Engagement**: Facilitates sharing and collaboration with a wider community.\n",
        "\n",
        "Connecting this notebook with the Hub is quite simple – all it requires is your Hub authentication token when asked. You can locate your Hub authentication token [here](https://huggingface.co/settings/tokens).\n",
        "\n",
        "**IMPORTANT**\n",
        "Make sure that your Hugging Face token was created with role `write` because the notebook will be saving the model to the Hugging Face repository.\n",
        "\n",
        "If you need instructions on how to create a Hugging Face tocken, please check out the [notebook for Chapter 3](https://github.com/PacktPublishing/Learn-OpenAI-Whisper/blob/main/Chapter03/LOAIW_ch03_working_with_audio_data_via_Hugging_Face.ipynb).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b045a39e-2a3e-4153-bdb5-281500bcd348",
      "metadata": {
        "id": "b045a39e-2a3e-4153-bdb5-281500bcd348"
      },
      "outputs": [],
      "source": [
        "from huggingface_hub import notebook_login\n",
        "\n",
        "notebook_login()"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "**IMPORTANT:**\n",
        "The output of `whoami90` shoud show `'role': 'write'` in your `auth` keys.. If it does not, then you might need to create a new token with role `write`, rerun the previous cell to execute `notebook_login()`, and load the new token.  \n",
        "\n",
        "<img src=\"https://github.com/PacktPublishing/Learn-OpenAI-Whisper/raw/main/Chapter04/hugging_face_whoami.png\">\n"
      ],
      "metadata": {
        "id": "5NgaNa3RjfNU"
      },
      "id": "5NgaNa3RjfNU"
    },
    {
      "cell_type": "code",
      "source": [
        "from huggingface_hub import whoami\n",
        "\n",
        "whoami()\n",
        "# you should see something like {'type': 'user',  'id': '...',  'name': 'Wauplin', ...}"
      ],
      "metadata": {
        "id": "t-Cj-HImlUSd"
      },
      "execution_count": null,
      "outputs": [],
      "id": "t-Cj-HImlUSd"
    },
    {
      "cell_type": "markdown",
      "id": "b219c9dd-39b6-4a95-b2a1-3f547a1e7bc0",
      "metadata": {
        "id": "b219c9dd-39b6-4a95-b2a1-3f547a1e7bc0"
      },
      "source": [
        "## Load dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "674429c5-0ab4-4adf-975b-621bb69eca38",
      "metadata": {
        "id": "674429c5-0ab4-4adf-975b-621bb69eca38"
      },
      "source": [
        "Utilizing 🤗 Datasets makes the task of downloading and preparing data remarkably straightforward.\n",
        "You can obtain and ready the Common Voice splits with a mere single line of code.\n",
        "\n",
        "Initially, confirm that you've agreed to the usage terms on the Hugging Face Hub: [mozilla-foundation/common_voice_11_0](https://huggingface.co/datasets/mozilla-foundation/common_voice_11_0). After agreeing to the terms, you gain complete access to the dataset, enabling you to download the data to your local environment.\n",
        "\n",
        "Given the scarcity of resources for Hindi, we plan to merge the `train` and `validation` splits, resulting in around 8 hours of training material. The 4 hours of `test` data will serve as our separate test set:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a2787582-554f-44ce-9f38-4180a5ed6b44",
      "metadata": {
        "id": "a2787582-554f-44ce-9f38-4180a5ed6b44"
      },
      "outputs": [],
      "source": [
        "from datasets import load_dataset, DatasetDict\n",
        "\n",
        "common_voice = DatasetDict()\n",
        "\n",
        "common_voice[\"train\"] = load_dataset(\"mozilla-foundation/common_voice_11_0\", \"hi\", split=\"train+validation\", use_auth_token=True)\n",
        "common_voice[\"test\"] = load_dataset(\"mozilla-foundation/common_voice_11_0\", \"hi\", split=\"test\", use_auth_token=True)\n",
        "\n",
        "print(common_voice)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d5c7c3d6-7197-41e7-a088-49b753c1681f",
      "metadata": {
        "id": "d5c7c3d6-7197-41e7-a088-49b753c1681f"
      },
      "source": [
        "The majority of ASR datasets offer just the input audio clips (`audio`) and their transcribed text (`sentence`). However, Common Voice includes extra metadata, like `accent` and `locale`, which aren't necessary for ASR tasks. To maintain the notebook's broad applicability, we focus solely on the input audio and the transcribed text for fine-tuning, omitting the extra metadata:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "20ba635d-518c-47ac-97ee-3cad25f1e0ce",
      "metadata": {
        "id": "20ba635d-518c-47ac-97ee-3cad25f1e0ce"
      },
      "outputs": [],
      "source": [
        "common_voice = common_voice.remove_columns([\"accent\", \"age\", \"client_id\", \"down_votes\", \"gender\", \"locale\", \"path\", \"segment\", \"up_votes\"])\n",
        "\n",
        "print(common_voice)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "2d63b2d2-f68a-4d74-b7f1-5127f6d16605",
      "metadata": {
        "id": "2d63b2d2-f68a-4d74-b7f1-5127f6d16605"
      },
      "source": [
        "## Prepare feature extractor, tokenizer and data"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "601c3099-1026-439e-93e2-5635b3ba5a73",
      "metadata": {
        "id": "601c3099-1026-439e-93e2-5635b3ba5a73"
      },
      "source": [
        "The ASR process can be broken down into three main components:\n",
        "1. A feature extractor that prepares the raw audio inputs\n",
        "2. The model that executes the sequence-to-sequence transformation\n",
        "3. A tokenizer that converts the model's outputs into text format\n",
        "\n",
        "Within 🤗 Transformers, the Whisper model comes with its own feature extractor and tokenizer, named [WhisperFeatureExtractor](https://huggingface.co/docs/transformers/main/model_doc/whisper#transformers.WhisperFeatureExtractor) and [WhisperTokenizer](https://huggingface.co/docs/transformers/main/model_doc/whisper#transformers.WhisperTokenizer), respectively.\n",
        "\n",
        "Let's dive into the specifics of configuring both the feature extractor and tokenizer step by step!"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "560332eb-3558-41a1-b500-e83a9f695f84",
      "metadata": {
        "id": "560332eb-3558-41a1-b500-e83a9f695f84"
      },
      "source": [
        "### Load WhisperFeatureExtractor"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "32ec8068-0bd7-412d-b662-0edb9d1e7365",
      "metadata": {
        "id": "32ec8068-0bd7-412d-b662-0edb9d1e7365"
      },
      "source": [
        "The Whisper feature extractor conducts two key tasks:\n",
        "1. Adjusts the length of the audio inputs to 30 seconds: it extends any audio shorter than 30 seconds with silence (zeros) to reach 30 seconds, and it cuts down any audio longer than 30 seconds to the same duration.\n",
        "2. Transforms the audio inputs into _log-Mel spectrogram_ input features, which visually represent the audio and match the input format the Whisper model anticipates."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "589d9ec1-d12b-4b64-93f7-04c63997da19",
      "metadata": {
        "id": "589d9ec1-d12b-4b64-93f7-04c63997da19"
      },
      "source": [
        "<figure>\n",
        "<img src=\"https://raw.githubusercontent.com/sanchit-gandhi/notebooks/main/spectrogram.jpg\" alt=\"Trulli\" style=\"width:100%\">\n",
        "<figcaption align = \"center\"><b>Figure 2:</b> Conversion of sampled audio array to log-Mel spectrogram.\n",
        "Left: sampled 1-dimensional audio signal. Right: corresponding log-Mel spectrogram. Figure source:\n",
        "<a href=\"https://ai.googleblog.com/2019/04/specaugment-new-data-augmentation.html\">Google SpecAugment Blog</a>.\n",
        "</figcaption>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b2ef54d5-b946-4c1d-9fdc-adc5d01b46aa",
      "metadata": {
        "id": "b2ef54d5-b946-4c1d-9fdc-adc5d01b46aa"
      },
      "source": [
        "We'll load the feature extractor from the pre-trained checkpoint with the default values:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "bc77d7bb-f9e2-47f5-b663-30f7a4321ce5",
      "metadata": {
        "id": "bc77d7bb-f9e2-47f5-b663-30f7a4321ce5"
      },
      "outputs": [],
      "source": [
        "from transformers import WhisperFeatureExtractor\n",
        "\n",
        "feature_extractor = WhisperFeatureExtractor.from_pretrained(\"openai/whisper-small\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "93748af7-b917-4ecf-a0c8-7d89077ff9cb",
      "metadata": {
        "id": "93748af7-b917-4ecf-a0c8-7d89077ff9cb"
      },
      "source": [
        "### Load WhisperTokenizer"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "2bc82609-a9fb-447a-a2af-99597c864029",
      "metadata": {
        "id": "2bc82609-a9fb-447a-a2af-99597c864029"
      },
      "source": [
        "The Whisper model generates a series of _token ids_. Each token id is then translated into its respective text string by the tokenizer. For the Hindi language, we are able to utilize the pre-trained tokenizer directly for fine-tuning without needing any adjustments. All that's required is to define the target language and the task at hand. This instruction guides the tokenizer to prepend the language and task tokens at the beginning of the encoded label sequences:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c7b07f9b-ae0e-4f89-98f0-0c50d432eab6",
      "metadata": {
        "id": "c7b07f9b-ae0e-4f89-98f0-0c50d432eab6"
      },
      "outputs": [],
      "source": [
        "from transformers import WhisperTokenizer\n",
        "\n",
        "tokenizer = WhisperTokenizer.from_pretrained(\"openai/whisper-small\", language=\"Hindi\", task=\"transcribe\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d2ef23f3-f4a8-483a-a2dc-080a7496cb1b",
      "metadata": {
        "id": "d2ef23f3-f4a8-483a-a2dc-080a7496cb1b"
      },
      "source": [
        "### Combine to create a WhisperProcessor"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5ff67654-5a29-4bb8-a69d-0228946c6f8d",
      "metadata": {
        "id": "5ff67654-5a29-4bb8-a69d-0228946c6f8d"
      },
      "source": [
        "To make the use of the feature extractor and tokenizer more straightforward, we can bundle them together into a unified `WhisperProcessor` class. This processor amalgamates functionalities from both the `WhisperFeatureExtractor` and `WhisperTokenizer`, allowing it to be applied to both the audio inputs and model outputs as needed. By adopting this approach, we simplify the components we need to manage during training to just two items: the `processor` and the `model`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "77d9f0c5-8607-4642-a8ac-c3ab2e223ea6",
      "metadata": {
        "id": "77d9f0c5-8607-4642-a8ac-c3ab2e223ea6"
      },
      "outputs": [],
      "source": [
        "from transformers import WhisperProcessor\n",
        "\n",
        "processor = WhisperProcessor.from_pretrained(\"openai/whisper-small\", language=\"Hindi\", task=\"transcribe\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "381acd09-0b0f-4d04-9eb3-f028ac0e5f2c",
      "metadata": {
        "id": "381acd09-0b0f-4d04-9eb3-f028ac0e5f2c"
      },
      "source": [
        "### Prepare data"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9649bf01-2e8a-45e5-8fca-441c13637b8f",
      "metadata": {
        "id": "9649bf01-2e8a-45e5-8fca-441c13637b8f"
      },
      "source": [
        "Let's print the first example of the Common Voice dataset to see\n",
        "what form the data is in:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6e6b0ec5-0c94-4e2c-ae24-c791be1b2255",
      "metadata": {
        "id": "6e6b0ec5-0c94-4e2c-ae24-c791be1b2255"
      },
      "outputs": [],
      "source": [
        "print(common_voice[\"train\"][0])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5a679f05-063d-41b3-9b58-4fc9c6ccf4fd",
      "metadata": {
        "id": "5a679f05-063d-41b3-9b58-4fc9c6ccf4fd"
      },
      "source": [
        "Given that our input audio has a sampling rate of 48kHz, it's essential to _downsample_ it to 16kHz before feeding it into the Whisper feature extractor, since 16kHz is the sampling rate the Whisper model expects.\n",
        "\n",
        "We will adjust the audio inputs to the appropriate sampling rate using the dataset's [`cast_column`](https://huggingface.co/docs/datasets/package_reference/main_classes.html?highlight=cast_column#datasets.DatasetDict.cast_column) function. This process doesn't alter the audio files directly. Instead, it instructs the `datasets` library to resample the audio samples _on the fly_ upon their initial access:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "f12e2e57-156f-417b-8cfb-69221cc198e8",
      "metadata": {
        "id": "f12e2e57-156f-417b-8cfb-69221cc198e8"
      },
      "outputs": [],
      "source": [
        "from datasets import Audio\n",
        "\n",
        "common_voice = common_voice.cast_column(\"audio\", Audio(sampling_rate=16000))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "00382a3e-abec-4cdd-a54c-d1aaa3ea4707",
      "metadata": {
        "id": "00382a3e-abec-4cdd-a54c-d1aaa3ea4707"
      },
      "source": [
        "Re-loading the first audio sample in the Common Voice dataset will resample\n",
        "it to the desired sampling rate:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "87122d71-289a-466a-afcf-fa354b18946b",
      "metadata": {
        "id": "87122d71-289a-466a-afcf-fa354b18946b"
      },
      "outputs": [],
      "source": [
        "print(common_voice[\"train\"][0])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "91edc72d-08f8-4f01-899d-74e65ce441fc",
      "metadata": {
        "id": "91edc72d-08f8-4f01-899d-74e65ce441fc"
      },
      "source": [
        "Now, let's craft a function to get our data set for the model:\n",
        "1. By invoking `batch[\"audio\"]`, we load and resample the audio data. As mentioned before, 🤗 Datasets handles the resampling as needed automatically.\n",
        "2. The feature extractor is utilized to derive the log-Mel spectrogram input features from our one-dimensional audio array.\n",
        "3. The transcriptions are converted into label ids by employing the tokenizer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6525c478-8962-4394-a1c4-103c54cce170",
      "metadata": {
        "id": "6525c478-8962-4394-a1c4-103c54cce170"
      },
      "outputs": [],
      "source": [
        "def prepare_dataset(batch):\n",
        "    # load and resample audio data from 48 to 16kHz\n",
        "    audio = batch[\"audio\"]\n",
        "\n",
        "    # compute log-Mel input features from input audio array\n",
        "    batch[\"input_features\"] = feature_extractor(audio[\"array\"], sampling_rate=audio[\"sampling_rate\"]).input_features[0]\n",
        "\n",
        "    # encode target text to label ids\n",
        "    batch[\"labels\"] = tokenizer(batch[\"sentence\"]).input_ids\n",
        "    return batch"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "70b319fb-2439-4ef6-a70d-a47bf41c4a13",
      "metadata": {
        "id": "70b319fb-2439-4ef6-a70d-a47bf41c4a13"
      },
      "source": [
        "We're able to implement the data preparation function across all training examples by utilizing the dataset's `.map` method. The `num_proc` parameter determines the number of CPU cores employed. Configuring `num_proc` to be greater than 1 activates multiprocessing. Should the `.map` method experience delays while using multiprocessing, adjusting `num_proc` to 1 will ensure the dataset is processed in a sequential manner."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7b73ab39-ffaf-4b9e-86e5-782963c6134b",
      "metadata": {
        "id": "7b73ab39-ffaf-4b9e-86e5-782963c6134b"
      },
      "outputs": [],
      "source": [
        "common_voice = common_voice.map(prepare_dataset, remove_columns=common_voice.column_names[\"train\"], num_proc=2)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "263a5a58-0239-4a25-b0df-c625fc9c5810",
      "metadata": {
        "id": "263a5a58-0239-4a25-b0df-c625fc9c5810"
      },
      "source": [
        "## Training and evaluation"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a693e768-c5a6-453f-89a1-b601dcf7daf7",
      "metadata": {
        "id": "a693e768-c5a6-453f-89a1-b601dcf7daf7"
      },
      "source": [
        "With our data now prepared, we're set to embark on setting up the training pipeline. The [🤗 Trainer](https://huggingface.co/transformers/master/main_classes/trainer.html?highlight=trainer) simplifies much of this process for us. Our tasks include:\n",
        "\n",
        "- Establishing a data collator: This function gathers our processed data and organizes it into PyTorch tensors, making them ready for the model.\n",
        "\n",
        "- Choosing evaluation metrics: For evaluation purposes, the model's performance will be assessed using the [word error rate (WER)](https://huggingface.co/metrics/wer) metric. We'll need a `compute_metrics` function that will manage this evaluation.\n",
        "\n",
        "- Loading a pre-trained checkpoint: A pre-trained checkpoint is loaded and set up for the training phase.\n",
        "\n",
        "- Configuring the training setup: The 🤗 Trainer relies on this configuration to outline the training protocol.\n",
        "\n",
        "After fine-tuning the model, we'll assess its performance on the test dataset to ensure it's been properly trained for speech transcription in Hindi."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8d230e6d-624c-400a-bbf5-fa660881df25",
      "metadata": {
        "id": "8d230e6d-624c-400a-bbf5-fa660881df25"
      },
      "source": [
        "### Define a data collator"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "04def221-0637-4a69-b242-d3f0c1d0ee78",
      "metadata": {
        "id": "04def221-0637-4a69-b242-d3f0c1d0ee78"
      },
      "source": [
        "The data collator for a sequence-to-sequence speech model distinctively processes `input_features` and `labels` separately: `input_features` are managed with the feature extractor while `labels` are dealt with by the tokenizer.\n",
        "\n",
        "`Input_features` get padded to 30 seconds and transformed into a log-Mel spectrogram of a uniform size through the feature extractor's intervention. The next step involves converting these `input_features` into batched PyTorch tensors, accomplished using the feature extractor's `.pad` method and setting `return_tensors=pt`.\n",
        "\n",
        "On the flip side, `labels` remain un-padded. The initial action is to pad these sequences to the batch's maximal length via the tokenizer's `.pad` method. Subsequently, padding tokens are substituted with `-100` to ensure they are excluded from loss calculations. Additionally, we remove the BOS token from the start of the label sequence since it will be added later in the training phase.\n",
        "\n",
        "The `WhisperProcessor` we previously established can be utilized to conduct both feature extraction and tokenization tasks:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8326221e-ec13-4731-bb4e-51e5fc1486c5",
      "metadata": {
        "id": "8326221e-ec13-4731-bb4e-51e5fc1486c5"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "\n",
        "from dataclasses import dataclass\n",
        "from typing import Any, Dict, List, Union\n",
        "\n",
        "@dataclass\n",
        "class DataCollatorSpeechSeq2SeqWithPadding:\n",
        "    processor: Any\n",
        "\n",
        "    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:\n",
        "        # split inputs and labels since they have to be of different lengths and need different padding methods\n",
        "        # first treat the audio inputs by simply returning torch tensors\n",
        "        input_features = [{\"input_features\": feature[\"input_features\"]} for feature in features]\n",
        "        batch = self.processor.feature_extractor.pad(input_features, return_tensors=\"pt\")\n",
        "\n",
        "        # get the tokenized label sequences\n",
        "        label_features = [{\"input_ids\": feature[\"labels\"]} for feature in features]\n",
        "        # pad the labels to max length\n",
        "        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors=\"pt\")\n",
        "\n",
        "        # replace padding with -100 to ignore loss correctly\n",
        "        labels = labels_batch[\"input_ids\"].masked_fill(labels_batch.attention_mask.ne(1), -100)\n",
        "\n",
        "        # if bos token is appended in previous tokenization step,\n",
        "        # cut bos token here as it's append later anyways\n",
        "        if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():\n",
        "            labels = labels[:, 1:]\n",
        "\n",
        "        batch[\"labels\"] = labels\n",
        "\n",
        "        return batch"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "3cae7dbf-8a50-456e-a3a8-7fd005390f86",
      "metadata": {
        "id": "3cae7dbf-8a50-456e-a3a8-7fd005390f86"
      },
      "source": [
        "Let's initialise the data collator we've just defined:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "fc834702-c0d3-4a96-b101-7b87be32bf42",
      "metadata": {
        "id": "fc834702-c0d3-4a96-b101-7b87be32bf42"
      },
      "outputs": [],
      "source": [
        "data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(common_voice)"
      ],
      "metadata": {
        "id": "zaDrQSTcnX48"
      },
      "id": "zaDrQSTcnX48",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "id": "d62bb2ab-750a-45e7-82e9-61d6f4805698",
      "metadata": {
        "id": "d62bb2ab-750a-45e7-82e9-61d6f4805698"
      },
      "source": [
        "### Evaluation metrics"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "66fee1a7-a44c-461e-b047-c3917221572e",
      "metadata": {
        "id": "66fee1a7-a44c-461e-b047-c3917221572e"
      },
      "source": [
        "We'll employ the word error rate (WER) metric, widely considered the standard for evaluating ASR systems. For additional details, you can consult the WER [documentation](https://huggingface.co/metrics/wer). The WER metric will be sourced from 🤗 Evaluate:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b22b4011-f31f-4b57-b684-c52332f92890",
      "metadata": {
        "id": "b22b4011-f31f-4b57-b684-c52332f92890"
      },
      "outputs": [],
      "source": [
        "import evaluate\n",
        "\n",
        "metric = evaluate.load(\"wer\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "4f32cab6-31f0-4cb9-af4c-40ba0f5fc508",
      "metadata": {
        "id": "4f32cab6-31f0-4cb9-af4c-40ba0f5fc508"
      },
      "source": [
        "We need to craft a function that processes our model predictions and calculates the WER metric. This function, named `compute_metrics`, initially substitutes `-100` with the `pad_token_id` in the `label_ids`, reversing the adjustment made in the data collator to accurately exclude padded tokens from the loss calculation. Subsequently, it translates the predicted and label ids into strings. Ultimately, it determines the WER by comparing the predictions with the reference labels:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "23959a70-22d0-4ffe-9fa1-72b61e75bb52",
      "metadata": {
        "id": "23959a70-22d0-4ffe-9fa1-72b61e75bb52"
      },
      "outputs": [],
      "source": [
        "def compute_metrics(pred):\n",
        "    pred_ids = pred.predictions\n",
        "    label_ids = pred.label_ids\n",
        "\n",
        "    # replace -100 with the pad_token_id\n",
        "    label_ids[label_ids == -100] = tokenizer.pad_token_id\n",
        "\n",
        "    # we do not want to group tokens when computing the metrics\n",
        "    pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)\n",
        "    label_str = tokenizer.batch_decode(label_ids, skip_special_tokens=True)\n",
        "\n",
        "    wer = 100 * metric.compute(predictions=pred_str, references=label_str)\n",
        "\n",
        "    return {\"wer\": wer}"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "daf2a825-6d9f-4a23-b145-c37c0039075b",
      "metadata": {
        "id": "daf2a825-6d9f-4a23-b145-c37c0039075b"
      },
      "source": [
        "### Load a pre-trained checkpoint"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "437a97fa-4864-476b-8abc-f28b8166cfa5",
      "metadata": {
        "id": "437a97fa-4864-476b-8abc-f28b8166cfa5"
      },
      "source": [
        "Now, we'll proceed to load the pre-trained Whisper `small` checkpoint, which is straightforward with the use of 🤗 Transformers!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5a10cc4b-07ec-4ebd-ac1d-7c601023594f",
      "metadata": {
        "id": "5a10cc4b-07ec-4ebd-ac1d-7c601023594f"
      },
      "outputs": [],
      "source": [
        "from transformers import WhisperForConditionalGeneration\n",
        "\n",
        "model = WhisperForConditionalGeneration.from_pretrained(\"openai/whisper-small\")\n",
        "model.generation_config.language = \"hi\" # Define the language of your choice here"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a15ead5f-2277-4a39-937b-585c2497b2df",
      "metadata": {
        "id": "a15ead5f-2277-4a39-937b-585c2497b2df"
      },
      "source": [
        "Adjust generation parameters - no tokens are predetermined as decoder outputs (refer to [`forced_decoder_ids`](https://huggingface.co/docs/transformers/main_classes/text_generation#transformers.generation_utils.GenerationMixin.generate.forced_decoder_ids)), and no tokens are excluded during the generation process (see [`suppress_tokens`](https://huggingface.co/docs/transformers/main_classes/text_generation#transformers.generation_utils.GenerationMixin.generate.suppress_tokens)):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "62038ba3-88ed-4fce-84db-338f50dcd04f",
      "metadata": {
        "id": "62038ba3-88ed-4fce-84db-338f50dcd04f"
      },
      "outputs": [],
      "source": [
        "model.config.forced_decoder_ids = None\n",
        "model.config.suppress_tokens = []"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "2178dea4-80ca-47b6-b6ea-ba1915c90c06",
      "metadata": {
        "id": "2178dea4-80ca-47b6-b6ea-ba1915c90c06"
      },
      "source": [
        "### Define the training configuration"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c21af1e9-0188-4134-ac82-defc7bdcc436",
      "metadata": {
        "id": "c21af1e9-0188-4134-ac82-defc7bdcc436"
      },
      "source": [
        "For the concluding step, we specify every parameter associated with the training process. To dive deeper into the specifics of training arguments, consult the Seq2SeqTrainingArguments [documentation](https://huggingface.co/docs/transformers/main_classes/trainer#transformers.Seq2SeqTrainingArguments)."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from huggingface_hub import whoami\n",
        "\n",
        "whoami()\n",
        "# you should see something like {'type': 'user',  'id': '...',  'name': 'Wauplin', ...}"
      ],
      "metadata": {
        "id": "M1gYmQVX28CO"
      },
      "id": "M1gYmQVX28CO",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0ae3e9af-97b7-4aa0-ae85-20b23b5bcb3a",
      "metadata": {
        "id": "0ae3e9af-97b7-4aa0-ae85-20b23b5bcb3a"
      },
      "outputs": [],
      "source": [
        "from transformers import Seq2SeqTrainingArguments\n",
        "\n",
        "training_args = Seq2SeqTrainingArguments(\n",
        "    output_dir=\"./whisper-small-hi\",  # change to a repo name of your choice\n",
        "    gradient_accumulation_steps=1,  # increase by 2x for every 2x decrease in batch size\n",
        "    learning_rate=5e-5,\n",
        "    # warmup_steps=500,\n",
        "    # max_steps=4000,\n",
        "    # warmup_steps=50,\n",
        "    # max_steps=400,\n",
        "    warmup_steps=10,\n",
        "    max_steps=100,\n",
        "    #num_train_epochs = 2.0,\n",
        "    gradient_checkpointing=True,\n",
        "    fp16=True,\n",
        "    evaluation_strategy=\"steps\",\n",
        "    per_device_train_batch_size=16,\n",
        "    per_device_eval_batch_size=10,\n",
        "    predict_with_generate=True,\n",
        "    generation_max_length=225,\n",
        "    # save_steps=1000,\n",
        "    # eval_steps=1000,\n",
        "    save_steps=100,\n",
        "    eval_steps=100,\n",
        "    logging_steps=10,\n",
        "    report_to=[\"tensorboard\"],\n",
        "    load_best_model_at_end=True,\n",
        "    metric_for_best_model=\"wer\",\n",
        "    greater_is_better=False,\n",
        "    hub_model_id = \"jbatista79/20240410-small-hindi\", # Change this to your own Hugging Face repo\n",
        "    # hub_token=\"\n",
        "    push_to_hub=True,\n",
        "    # push_to_hub=False,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b3a944d8-3112-4552-82a0-be25988b3857",
      "metadata": {
        "id": "b3a944d8-3112-4552-82a0-be25988b3857"
      },
      "source": [
        "**Note**: if one does not want to upload the model checkpoints to the Hub,\n",
        "set `push_to_hub=False`."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "bac29114-d226-4f54-97cf-8718c9f94e1e",
      "metadata": {
        "id": "bac29114-d226-4f54-97cf-8718c9f94e1e"
      },
      "source": [
        "We can pass the training arguments to the 🤗 Trainer, together with our model, dataset, data collator, and `compute_metrics` function:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d546d7fe-0543-479a-b708-2ebabec19493",
      "metadata": {
        "id": "d546d7fe-0543-479a-b708-2ebabec19493"
      },
      "outputs": [],
      "source": [
        "from transformers import Seq2SeqTrainer\n",
        "\n",
        "trainer = Seq2SeqTrainer(\n",
        "    args=training_args,\n",
        "    model=model,\n",
        "    train_dataset=common_voice[\"train\"],\n",
        "    eval_dataset=common_voice[\"test\"],\n",
        "    data_collator=data_collator,\n",
        "    compute_metrics=compute_metrics,\n",
        "    tokenizer=processor.feature_extractor,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Before initiating the training process, we'll save the processor object. Given that the processor is non-trainable, it will remain unchanged throughout the training period:"
      ],
      "metadata": {
        "id": "uOrRhDGtN5S4"
      },
      "id": "uOrRhDGtN5S4"
    },
    {
      "cell_type": "code",
      "source": [
        "processor.save_pretrained(training_args.output_dir)"
      ],
      "metadata": {
        "id": "-2zQwMfEOBJq"
      },
      "id": "-2zQwMfEOBJq",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "id": "7f404cf9-4345-468c-8196-4bd101d9bd51",
      "metadata": {
        "id": "7f404cf9-4345-468c-8196-4bd101d9bd51"
      },
      "source": [
        "### Training"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5e8b8d56-5a70-4f68-bd2e-f0752d0bd112",
      "metadata": {
        "id": "5e8b8d56-5a70-4f68-bd2e-f0752d0bd112"
      },
      "source": [
        "The training duration is estimated to be around 5-10 hours, contingent on the capabilities of your GPU or the one provided by this Google Colab. If you're utilizing this Google Colab for fine-tuning a Whisper model, it's crucial to ensure that the training process does not get halted due to periods of inactivity. To circumvent this issue, you can employ a straightforward solution by inserting the following script into the console of this browser tab (by right-clicking -> selecting _inspect_ -> navigating to the _Console_ tab -> pasting the code)."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "890a63ed-e87b-4e53-a35a-6ec1eca560af",
      "metadata": {
        "id": "890a63ed-e87b-4e53-a35a-6ec1eca560af"
      },
      "source": [
        "```javascript\n",
        "function ConnectButton(){\n",
        "    console.log(\"Connect pushed\");\n",
        "    document.querySelector(\"#top-toolbar > colab-connect-button\").shadowRoot.querySelector(\"#connect\").click()\n",
        "}\n",
        "setInterval(ConnectButton, 60000);\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5a55168b-2f46-4678-afa0-ff22257ec06d",
      "metadata": {
        "id": "5a55168b-2f46-4678-afa0-ff22257ec06d"
      },
      "source": [
        "The maximum GPU memory usage for this training setup is roughly 15.8GB. If the GPU assigned to your Google Colab session is not sufficient, you may run into a CUDA `\"out-of-memory\"` error upon initiating the training. Should this occur, a viable strategy is to gradually decrease the `per_device_train_batch_size` by halves and utilize [`gradient_accumulation_steps`](https://huggingface.co/docs/transformers/main_classes/trainer#transformers.Seq2SeqTrainingArguments.gradient_accumulation_steps) as a balancing measure.\n",
        "\n",
        "To begin the training process, just run:"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "gpu_info = !nvidia-smi\n",
        "gpu_info = '\\n'.join(gpu_info)\n",
        "if gpu_info.find('failed') >= 0:\n",
        "  print('Not connected to a GPU')\n",
        "else:\n",
        "  print(gpu_info)"
      ],
      "metadata": {
        "id": "PFnWSLUQocDL"
      },
      "id": "PFnWSLUQocDL",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import locale\n",
        "locale.getpreferredencoding = lambda: \"UTF-8\"\n",
        "locale.getpreferredencoding()\n"
      ],
      "metadata": {
        "id": "Vyt-Ene4yh-J"
      },
      "id": "Vyt-Ene4yh-J",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ee8b7b8e-1c9a-4d77-9137-1778a629e6de",
      "metadata": {
        "id": "ee8b7b8e-1c9a-4d77-9137-1778a629e6de"
      },
      "outputs": [],
      "source": [
        "trainer.train()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "34d4360d-5721-426e-b6ac-178f833fedeb",
      "metadata": {
        "id": "34d4360d-5721-426e-b6ac-178f833fedeb"
      },
      "source": [
        "## Building a demo"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [],
      "metadata": {
        "id": "xX3NOiTJVbqZ"
      },
      "id": "xX3NOiTJVbqZ"
    },
    {
      "cell_type": "markdown",
      "id": "e65489b7-18d1-447c-ba69-cd28dd80dad9",
      "metadata": {
        "id": "e65489b7-18d1-447c-ba69-cd28dd80dad9"
      },
      "source": [
        "Having fine-tuned our model, we're now ready to demonstrate its ASR prowess through a demo. We'll employ the 🤗 Transformers `pipeline` function, which seamlessly manages the entire ASR process, from preprocessing the audio inputs to decoding the model's predictions.\n",
        "\n",
        "Executing the code below will create a Gradio demo. This demo allows us to capture speech using our computer's microphone and submit it to our fine-tuned Whisper model for text transcription:"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install git+https://github.com/huggingface/transformers\n",
        "!pip install gradio"
      ],
      "metadata": {
        "id": "H9WteS-JW_1T"
      },
      "id": "H9WteS-JW_1T",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e0ace3aa-1ef3-45cb-933f-6ddca037c5aa",
      "metadata": {
        "id": "e0ace3aa-1ef3-45cb-933f-6ddca037c5aa"
      },
      "outputs": [],
      "source": [
        "from transformers import pipeline\n",
        "import gradio as gr\n",
        "\n",
        "pipe = pipeline(model=\"jbatista79/20240410-small-hindi\")  # change to \"your-username/the-name-you-picked\"\n",
        "\n",
        "def transcribe(audio):\n",
        "    text = pipe(audio)[\"text\"]\n",
        "    return text\n",
        "\n",
        "iface = gr.Interface(\n",
        "    fn=transcribe,\n",
        "    # inputs=gr.Audio(source=\"microphone\", type=\"filepath\"),\n",
        "    inputs=gr.Audio(sources=[\"microphone\"], type=\"filepath\"),\n",
        "    outputs=\"text\",\n",
        "    title=\"Whisper Small Hindi\",\n",
        "    description=\"Realtime demo for Hindi speech recognition using a fine-tuned Whisper small model.\",\n",
        ")\n",
        "\n",
        "iface.launch()"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Gratitude\n",
        "\n",
        "Many thanks to Naval Katoch for his valuable technical insights. For more details on the Whisper model, the Common Voice dataset and the theory behind fine-tuning, check out the pioneering work by Hugging Face's Sanchit Gandhi: [blog post](https://huggingface.co/blog/fine-tune-whisper)."
      ],
      "metadata": {
        "id": "1_qasqaAKHPR"
      },
      "id": "1_qasqaAKHPR"
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.9"
    },
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 5
}